{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "5obtia7wRSae",
        "cJtau0HASK4K",
        "iO7YxIbTRZ-b"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install faker"
      ],
      "metadata": {
        "id": "yD06d33fOK-y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configure and Load the NER Model\n"
      ],
      "metadata": {
        "id": "cJtau0HASK4K"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install and import required python libraries"
      ],
      "metadata": {
        "id": "Ly1f1P-l9ri8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Elastic uses the [eland python library](https://github.com/elastic/eland) to download modesl from Hugging Face hub and load them into elasticsearch"
      ],
      "metadata": {
        "id": "MJAb_8zlPFhQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rUedSzQW9FIF"
      },
      "outputs": [],
      "source": [
        "!pip install eland elasticsearch transformers sentence_transformers torch==1.11"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pathlib import Path\n",
        "from eland.ml.pytorch import PyTorchModel\n",
        "from eland.ml.pytorch.transformers import TransformerModel\n",
        "from elasticsearch import Elasticsearch\n",
        "from elasticsearch.client import MlClient"
      ],
      "metadata": {
        "id": "wyUZXUi4RWWL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Configure elasticsearch authentication.\n",
        "The recommended authentication approach is using the [Elastic Cloud ID](https://www.elastic.co/guide/en/cloud/current/ec-cloud-id.html) and a [cluster level API key](https://www.elastic.co/guide/en/kibana/current/api-keys.html)\n",
        "\n",
        "You can use any method you wish to set the required credentials. We are using getpass in this example to prompt for credentials to avoide storing them in github."
      ],
      "metadata": {
        "id": "r7nMIbHke37Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import getpass"
      ],
      "metadata": {
        "id": "Xsd2m7HoTCLm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "es_cloud_id = getpass.getpass('Enter Elastic Cloud ID:  ')\n",
        "es_api_key = getpass.getpass('Enter cluster API key:  ')"
      ],
      "metadata": {
        "id": "SSGgYHome69o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Connect to Elastic Cloud"
      ],
      "metadata": {
        "id": "jL4VDnVp96lf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "es = Elasticsearch(cloud_id=es_cloud_id,\n",
        "#                   api_key=(es_api_id, es_api_key)\n",
        "                   api_key=es_api_key\n",
        "                   )\n",
        "es.info() # should return cluster info"
      ],
      "metadata": {
        "id": "I8mVJkKmetXo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading a model from Hugging Face model hub"
      ],
      "metadata": {
        "id": "pOuF_1VnmVU-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The Elastic Stack machine learning features support transformer models that conform to the standard BERT model interface and use the WordPiece tokenization algorithm.\n",
        "\n",
        "A current list of supported architectures can be viewed in the\n",
        "[Elastic NLP Model Support Docs](https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-model-ref.html)\n"
      ],
      "metadata": {
        "id": "NwIItJyhoWeT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download an NER model from Hugging Face using the HF copy link\n",
        "\n",
        "For this example we will be using the [sentence-transformers/msmarco-MiniLM-L-12-v3](https://huggingface.co/sentence-transformers/msmarco-MiniLM-L-12-v3) model\n"
      ],
      "metadata": {
        "id": "10VvWJ87alld"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download the Model\n",
        "Here we specify the model id from Hugging Face. The easiest way to get this id is clicking the copy the model name icon next to the name on the model page.\n",
        "\n",
        "When calling `TransformerModel` you specify the HF model id and the task type. You can try specifying `auto` and eland will attempt to determine the correct type from info in the model config. This is not always possible so a list of specific `task_type` values can be viewed in the following code:\n",
        "[Supported values](https://github.com/elastic/eland/blob/15a300728876022b206161d71055c67b500a0192/eland/ml/pytorch/transformers.py#*L41*)\n",
        "\n",
        "_note and warnings about \"Some weights of the model checkpoint\" can be ignored_"
      ],
      "metadata": {
        "id": "uBMWHj-ZmtvE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "hf_model_id='dslim/bert-base-NER'\n",
        "tm = TransformerModel(model_id=hf_model_id, task_type=\"ner\")"
      ],
      "metadata": {
        "id": "zPV3oFsKiYFL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set and confirm the model ID\n",
        "To make the name compatible with elasticsearch, the '/' is replaced with '__'\n",
        "\n"
      ],
      "metadata": {
        "id": "sX-9dHuDmwgX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "es_model_id = tm.elasticsearch_model_id()\n",
        "es_model_id"
      ],
      "metadata": {
        "id": "XkIQBBCbdqvQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Export the model in a TorchScrpt representation which Elasticsearch uses"
      ],
      "metadata": {
        "id": "p0L2cfYwbIld"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tmp_path = \"models\"\n",
        "Path(tmp_path).mkdir(parents=True, exist_ok=True)\n",
        "model_path, config, vocab_path = tm.save(tmp_path)"
      ],
      "metadata": {
        "id": "GsSpvvP-nbCK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load the model into Elasticsearch\n",
        "Model should not already exist in elasticsearch"
      ],
      "metadata": {
        "id": "k1a_yNo6ba2E"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ptm = PyTorchModel(es, es_model_id)\n",
        "ptm.import_model(model_path=model_path, config_path=None, vocab_path=vocab_path, config=config)"
      ],
      "metadata": {
        "id": "Z4QD71Apnj4j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Starting the Model"
      ],
      "metadata": {
        "id": "4UYSzFp3vHdB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### View information about the model\n",
        "This is not required but can be handy to get a model overivew"
      ],
      "metadata": {
        "id": "wQwfozwznK4Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# List the in elasticsearch\n",
        "m = MlClient.get_trained_models(es, model_id=es_model_id)\n",
        "m.body"
      ],
      "metadata": {
        "id": "b4Wv8EJvpfZI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Deploy the model\n",
        "This will load the model on the ML nodes and start the process(es) making it available for the NLP task"
      ],
      "metadata": {
        "id": "oMGw3sk-pbaN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "s = MlClient.start_trained_model_deployment(es, model_id=es_model_id)\n",
        "s.body"
      ],
      "metadata": {
        "id": "w5muJ1rLqvUW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Verify the model started without issue"
      ],
      "metadata": {
        "id": "ZytlELrsnn_O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "stats = MlClient.get_trained_models_stats(es, model_id=es_model_id)\n",
        "stats.body['trained_model_stats'][0]['deployment_stats']['nodes'][0]['routing_state']"
      ],
      "metadata": {
        "id": "ZaQUUWe0Hxwz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create the ingest pipeline"
      ],
      "metadata": {
        "id": "iO7YxIbTRZ-b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "description= \"PII redacting ingest pipeline\"\n",
        "\n",
        "processors= [\n",
        "    {\n",
        "      \"set\": {\n",
        "        \"field\": \"redacted\",\n",
        "        \"value\": \"{{{message}}}\"\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "      \"inference\": {\n",
        "        \"model_id\": es_model_id,\n",
        "        \"field_map\": {\n",
        "          \"message\": \"text_field\"\n",
        "        }\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "      \"script\": {\n",
        "        \"lang\": \"painless\",\n",
        "        \"source\": \"\"\"String msg = ctx['message'];\n",
        "      for (item in ctx['ml']['inference']['entities']) {\n",
        "        msg = msg.replace(item['entity'], '<' + item['class_name'] + '>')\n",
        "      }\n",
        "      ctx['redacted']=msg\"\"\",\n",
        "        \"if\": \"return ctx['ml']['inference']['entities'].isEmpty() == false\",\n",
        "        \"tag\": \"ner_redact\",\n",
        "        \"description\": \"Redact NER entities\"\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "      \"redact\": {\n",
        "        \"field\": \"redacted\",\n",
        "        \"patterns\": [\n",
        "          \"%{PHONE:PHONE}\",\n",
        "          \"%{SSN:SSN}\"\n",
        "        ],\n",
        "        \"pattern_definitions\": {\n",
        "          \"SSN\": \"\"\"\\d{3}-?\\d{2}-?\\d{4}\"\"\",\n",
        "          \"PHONE\": \"\"\"\\d{3}-?\\d{3}-?\\d{4}\"\"\"\n",
        "        }\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "      \"remove\": {\n",
        "        \"field\": [\n",
        "          \"message\",\n",
        "          \"ml\"\n",
        "        ]\n",
        "      }\n",
        "    }\n",
        "  ]\n",
        "\n",
        "on_failure= [\n",
        "    {\n",
        "      \"set\": {\n",
        "        \"field\": \"failure\",\n",
        "        \"value\": \"pii_script-redact\"\n",
        "      }\n",
        "    }\n",
        "  ]\n",
        "\n",
        "\n",
        "\n",
        "response = es.ingest.put_pipeline(id=\"pii_redaction_pipeline_book\",\n",
        "                                  description=description,\n",
        "                                  processors=processors,\n",
        "                                  on_failure=on_failure\n",
        ")\n",
        "\n",
        "\n",
        "# Print the response\n",
        "print(response)"
      ],
      "metadata": {
        "id": "DsUkWlNkNsSB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configure the index template\n",
        "Match any indices created with the name pattern `pii_data*`\n",
        "Create the mappings for the fields\n",
        "Configure new data to use our `pii_redaction_pipeline` by default"
      ],
      "metadata": {
        "id": "YuORHcRMTd0B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "index_patterns =[\n",
        "    \"pii_data*\"\n",
        "  ]\n",
        "order = 1\n",
        "settings = {\n",
        "      \"number_of_shards\": 1,\n",
        "      \"number_of_replicas\": 1,\n",
        "      \"index.default_pipeline\": \"pii_redaction_pipeline_book\"\n",
        "    }\n",
        "mappings = {\n",
        "      \"properties\": {\n",
        "        \"message\": {\n",
        "          \"type\": \"text\"\n",
        "        },\n",
        "        \"status\": {\n",
        "          \"type\": \"keyword\"\n",
        "        },\n",
        "        \"redacted\": {\n",
        "          \"type\": \"text\"\n",
        "        }\n",
        "      }\n",
        "    }\n",
        "\n",
        "\n",
        "# Create the index template\n",
        "response = es.indices.put_template(name=\"pii_book_template\",\n",
        "                                   index_patterns=index_patterns,\n",
        "                                   order=order,\n",
        "                                   settings=settings,\n",
        "                                   mappings=mappings\n",
        "                                   )\n",
        "\n",
        "# Print the response\n",
        "print(response)\n"
      ],
      "metadata": {
        "id": "dD4Q_tHxTjYi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate fake PII data"
      ],
      "metadata": {
        "id": "5obtia7wRSae"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99PLiMoHNo75"
      },
      "outputs": [],
      "source": [
        "from faker import Faker\n",
        "import json\n",
        "from pprint import pprint\n",
        "\n",
        "# Create an instance of the Faker class\n",
        "fake = Faker()\n",
        "\n",
        "# Define a function to generate fake personal identification information\n",
        "def generate_fake_pii(num_records):\n",
        "\n",
        "  fake_data = []\n",
        "\n",
        "  for x in range(num_records):\n",
        "  # Generate fake personal identification information\n",
        "    fn = fake.first_name()\n",
        "    ln = fake.last_name()\n",
        "    pn = fake.phone_number()\n",
        "    sn = fake.ssn()\n",
        "    ai = fake.random_element(elements=('active', 'inactive'))\n",
        "\n",
        "    call_log = {\n",
        "        'message' : f'{fn} {ln} called in from {pn}. Their account number is {sn}',\n",
        "        'status' : ai\n",
        "        }\n",
        "\n",
        "    fake_data.append(call_log)\n",
        "  return fake_data\n",
        "\n",
        "# Generate fake personal identification information for N individuals\n",
        "num_records = 10 # Set the desired number of records\n",
        "fake_pii_data = generate_fake_pii(num_records)\n",
        "\n",
        "pprint(fake_pii_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Ingest Fake data into Elasticsearch\n",
        "Loading the data to the pii_data index will pass the data through the configured ingest pipeline redacting pii."
      ],
      "metadata": {
        "id": "cQKs5Y79TPsw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from elasticsearch import Elasticsearch, helpers\n",
        "\n",
        "# Define a function to generate an array of documents for Elasticsearch\n",
        "def generate_documents_array(fake_data, index_name):\n",
        "\n",
        "  documents = []\n",
        "\n",
        "  # Iterate over the fake data and format it as documents\n",
        "  for call in fake_data:\n",
        "    # Create a document with the _index and _source keys\n",
        "    document = {\n",
        "      '_index': index_name,\n",
        "      '_source': call\n",
        "    }\n",
        "\n",
        "    documents.append(document)\n",
        "\n",
        "  return documents\n",
        "\n",
        "\n",
        "# Generate the array of documents for Elasticsearch\n",
        "index_name = 'pii_data-book'\n",
        "documents_array = generate_documents_array(fake_pii_data, index_name)\n",
        "\n",
        "# Convert the bulk index body to a single string with newline separators\n",
        "print(\"Bulk request: \")\n",
        "print(documents_array)\n",
        "\n",
        "try:\n",
        "\tresponse = helpers.bulk(es, documents_array)\n",
        "\tprint (\"\\nRESPONSE:\", response)\n",
        "except Exception as e:\n",
        "\tprint(\"\\nERROR:\", e)"
      ],
      "metadata": {
        "id": "paxHUfguOu88"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = {\"match_all\":{}}\n",
        "\n",
        "response = es.search(\n",
        "    index=index_name,\n",
        "    query=query\n",
        "    )\n",
        "\n",
        "pprint(response['hits']['hits'])"
      ],
      "metadata": {
        "id": "1t7EKO5xFj3_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}
